{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Objax Basics Tutorial",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fRTfLE9T_hB_",
        "colab_type": "text"
      },
      "source": [
        "# Objax Basics Tutorial\n",
        "\n",
        "This tutorial introduces basic Objax concepts.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1mGUIErXGmZq",
        "colab_type": "text"
      },
      "source": [
        "## Prerequisites\n",
        "\n",
        "Objax is a machine learning library written in Python which works on top of [JAX](https://jax.readthedocs.io/en/latest/index.html#). Readers should therefore have some familiarity with following:\n",
        "\n",
        "* **Python**. If you're new to Python or need a refresher check [Python Tutorial](https://docs.python.org/3/tutorial/) or [Python Beginner's Guide](https://wiki.python.org/moin/BeginnersGuide).\n",
        "* **NumPy** is a library for mathematical computations. Many of JAX primitives are based on NumPy and have the same syntax. NumPy is also useful for data manipulation outside of JAX/Objax. [NumPy quickstart](https://numpy.org/doc/stable/user/quickstart.html) covers most of the needed topics. More information can be found on [NumPy documentation site](https://numpy.org/doc/stable/).\n",
        "* **JAX** can be described as NumPy with gradients which runs on accelerators (GPU and TPU). The [JAX quickstart](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html) covers most of the concepts needed to understand Objax."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "erb7-6AFGo6T",
        "colab_type": "text"
      },
      "source": [
        "## Installation and imports\n",
        "\n",
        "Let's first install Objax:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U2fVDSdK6RLk",
        "colab_type": "code",
        "colab": {},
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772836743,
          "user_tz": 420,
          "elapsed": 333,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        }
      },
      "source": [
        "%pip --quiet install objax"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6m1yKOUv6TUZ",
        "colab_type": "text"
      },
      "source": [
        "After Objax is installed, you can import all necessary modules:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vc4zUzUy_TiW",
        "colab_type": "code",
        "colab": {},
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772853174,
          "user_tz": 420,
          "elapsed": 16410,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        }
      },
      "source": [
        "import jax.numpy as jn\n",
        "import numpy as np\n",
        "import objax"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "62YieL6YGef-",
        "colab_type": "text"
      },
      "source": [
        "## Tensors\n",
        "\n",
        "Tensors are essentially multi-dimensional arrays. In JAX and Objax tensors can be placed on GPUs or TPUs to accelerate computations.\n",
        "\n",
        "Objax relies on the `jax.numpy.ndarray` primitive from JAX to represent tensors. In turn, this primitive has a very similar API to NumPy `ndarray`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9uSJ7f5Ho0K",
        "colab_type": "text"
      },
      "source": [
        "### Creating tensors\n",
        "\n",
        "Tensors creation is very similar to NumPy and is done in multiple ways:\n",
        "\n",
        "1. Provide explicit values to the tensor:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "awwnqbFhGiH-",
        "colab_type": "code",
        "colab": {
          "height": 105
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772853563,
          "user_tz": 420,
          "elapsed": 357,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "ba83d016-85ba-45ec-8b1f-913a37addb56"
      },
      "source": [
        "# Providing explicit values\n",
        "jn.array([[1.0, 2.0, 3.0],\n",
        "          [4.0, 5.0, 6.0]])"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[1., 2., 3.],\n",
              "             [4., 5., 6.]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1eHxxHd7F_q3",
        "colab_type": "text"
      },
      "source": [
        "2. From a NumPy array:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-FqWRj--GJsg",
        "colab_type": "code",
        "colab": {
          "height": 34
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772853887,
          "user_tz": 420,
          "elapsed": 304,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "c7d03aea-8772-4413-f719-31951b4868de"
      },
      "source": [
        "arr = np.array([1.0, 2.0, 3.0])\n",
        "jn.array(arr)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([1., 2., 3.], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1HyqxxnqG-bd",
        "colab_type": "text"
      },
      "source": [
        "3. From another JAX tensor:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XUrlrrY4zjH0",
        "colab_type": "code",
        "colab": {
          "height": 51
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772854204,
          "user_tz": 420,
          "elapsed": 308,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "c78b8966-38b2-4b69-ada4-81e2c77a3348"
      },
      "source": [
        "another_tensor = jn.array([[1.0, 2.0, 3.0],\n",
        "                           [4.0, 5.0, 6.0]])\n",
        "jn.array(another_tensor)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[1., 2., 3.],\n",
              "             [4., 5., 6.]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6G1PGpv1GIzW",
        "colab_type": "text"
      },
      "source": [
        "4. Using `ones` or `zeros`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w8osctRNF-zM",
        "colab_type": "code",
        "colab": {
          "height": 68
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772854543,
          "user_tz": 420,
          "elapsed": 332,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "386b3629-6f90-4a2e-ff7c-f8893a6e6aba"
      },
      "source": [
        "jn.ones((3, 4))"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[1., 1., 1., 1.],\n",
              "             [1., 1., 1., 1.],\n",
              "             [1., 1., 1., 1.]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GtaICIZ3G-SW",
        "colab_type": "code",
        "colab": {
          "height": 85
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772854874,
          "user_tz": 420,
          "elapsed": 314,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "fecc7e38-f8b3-4d4e-df68-b390d44d3e31"
      },
      "source": [
        "jn.zeros((4, 5))"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[0., 0., 0., 0., 0.],\n",
              "             [0., 0., 0., 0., 0.],\n",
              "             [0., 0., 0., 0., 0.],\n",
              "             [0., 0., 0., 0., 0.]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BnQGslSZHDPn",
        "colab_type": "text"
      },
      "source": [
        "5. As a result of a mathematical operation performed on other tensors:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4LZzXANjHMsx",
        "colab_type": "code",
        "colab": {
          "height": 51
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772855244,
          "user_tz": 420,
          "elapsed": 352,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "95457593-9cac-4233-fad1-fafcaad30001"
      },
      "source": [
        "t1 = jn.array([[1.0, 2.0, 3.0],\n",
        "               [4.0, 5.0, 6.0]])\n",
        "t2 = jn.ones(t1.shape) * 3\n",
        "t1 + t2"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DeviceArray([[4., 5., 6.],\n",
              "             [7., 8., 9.]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WrbRRgWXH8wg",
        "colab_type": "text"
      },
      "source": [
        "### Tensor Properties\n",
        "\n",
        "Similar to NumPy, one can explore various properties of tensors like shape, number of dimensions, or data type:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "slSfCg-cIAA-",
        "colab_type": "code",
        "colab": {
          "height": 68
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772855561,
          "user_tz": 420,
          "elapsed": 306,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "5b3be06c-9db4-4790-deda-41e87313725d"
      },
      "source": [
        "t = jn.array([[1.0, 2.0, 3.0],\n",
        "              [4.0, 5.0, 6.0]])\n",
        "print('Number of dimensions: ', t.ndim)\n",
        "print('Shape: ', t.shape)\n",
        "print('Data type: ', t.dtype)"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Number of dimensions:  2\n",
            "Shape:  (2, 3)\n",
            "Data type:  float32\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "or1yU-HWHQMa",
        "colab_type": "text"
      },
      "source": [
        "### Converting tensors to numpy array\n",
        "\n",
        "Objax/JAX tensors can be converted to NumPy arrays when needed to perform computations with NumPy:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "raIwjxOZHZnL",
        "colab_type": "code",
        "colab": {
          "height": 51
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772855898,
          "user_tz": 420,
          "elapsed": 318,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "8ae5592b-4a04-4209-c7b1-6313a6d5dc7a"
      },
      "source": [
        "np.array(t)"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[1., 2., 3.],\n",
              "       [4., 5., 6.]], dtype=float32)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bGXfS_QcI0mb",
        "colab_type": "text"
      },
      "source": [
        "### Tensors are immutable\n",
        "\n",
        "One big difference between JAX `ndarray` and NumPy `ndarray` is that JAX `ndarray` is immutable:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FILbMlAHIVuW",
        "colab_type": "code",
        "colab": {
          "height": 136
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772856242,
          "user_tz": 420,
          "elapsed": 318,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "567dd8d3-6bf1-427c-ef45-a78229a7a114"
      },
      "source": [
        "print('Original tensor t:\\n', t)\n",
        "\n",
        "try:\n",
        "    t[0, 0] = -5.0  # This line will fail\n",
        "except Exception as e:\n",
        "    print(f'Exception {e}')\n",
        "\n",
        "print('Tensor t after failed attempt to update:\\n', t)"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Original tensor t:\n",
            " [[1. 2. 3.]\n",
            " [4. 5. 6.]]\n",
            "Exception '<class 'jax.interpreters.xla.DeviceArray'>' object does not support item assignment. JAX arrays are immutable; perhaps you want jax.ops.index_update or jax.ops.index_add instead?\n",
            "Tensor t after failed attempt to update:\n",
            " [[1. 2. 3.]\n",
            " [4. 5. 6.]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "STMiM6ElIsW9",
        "colab_type": "text"
      },
      "source": [
        "Instead of updating an existing tensor, a new tensor should be created with updated elements. Updates of individual tensor elements is done using\n",
        "[index_update](https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.index_update.html#jax.ops.index_update), [index_add](https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.index_add.html#jax.ops.index_add) and some other JAX primitives:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nr3NZR6hIrne",
        "colab_type": "code",
        "colab": {
          "height": 170
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772856653,
          "user_tz": 420,
          "elapsed": 393,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "bdf23796-bc90-4026-f49d-44309e556b04"
      },
      "source": [
        "import jax.ops\n",
        "\n",
        "print('Original tensor t:\\n', t)\n",
        "new_t = jax.ops.index_update(t, jax.ops.index[0, 0], -5.0)\n",
        "print('Tensor t after update stays the same:\\n', t)\n",
        "print('Tensor new_t has updated value:\\n', new_t)"
      ],
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Original tensor t:\n",
            " [[1. 2. 3.]\n",
            " [4. 5. 6.]]\n",
            "Tensor t after update stays the same:\n",
            " [[1. 2. 3.]\n",
            " [4. 5. 6.]]\n",
            "Tensor new_t has updated value:\n",
            " [[-5.  2.  3.]\n",
            " [ 4.  5.  6.]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r8FI0dSKI3t0",
        "colab_type": "text"
      },
      "source": [
        "More details about per-element updates of tensors can be found in [JAX documentation](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-In-Place-Updates).\n",
        "\n",
        "In practice, most mathematical operations cover tensors as a whole while manual per-element updates are rarely needed."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AM9M0AtjGiyS",
        "colab_type": "text"
      },
      "source": [
        "## Random numbers\n",
        "\n",
        "It's very easy to generate random tensors in Objax:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "98J3YqqgGwGv",
        "colab_type": "code",
        "colab": {
          "height": 68
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772857591,
          "user_tz": 420,
          "elapsed": 915,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "b0b78c27-fd0a-4db9-fdab-4863f6334f23"
      },
      "source": [
        "x = objax.random.normal((3, 4))\n",
        "print(x)"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[[-0.1441347   0.89507747 -0.46038115  0.10503326]\n",
            " [-0.7460886   0.89681065  0.38794124  0.11750659]\n",
            " [ 1.0659382   0.22656879 -2.548792    1.9700414 ]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f7_kM1BECPz3",
        "colab_type": "text"
      },
      "source": [
        "There are multiple primitives for doing so:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tL1Gm6-smGeZ",
        "colab_type": "code",
        "colab": {
          "height": 85
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772859147,
          "user_tz": 420,
          "elapsed": 1530,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "5fd2f555-7db9-40f8-fbfe-eeaf0089a749"
      },
      "source": [
        "print('Random integers:', objax.random.randint((4,), low=0, high=10))\n",
        "print('Random normal:', objax.random.normal((4,), mean=1.0, stddev=2.0))\n",
        "print('Random truncated normal: ', objax.random.truncated_normal((4,), stddev=2.0))\n",
        "print('Random uniform: ', objax.random.uniform((4,)))"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Random integers: [1 3 4 7]\n",
            "Random normal: [-3.4102912  -1.8277478  -0.02106905  0.65338284]\n",
            "Random truncated normal:  [ 0.78602946  2.3004575  -0.22719319  0.22819921]\n",
            "Random uniform:  [0.19456518 0.4642099  0.94732213 0.57298625]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b7cNGncYGw7y",
        "colab_type": "text"
      },
      "source": [
        "## Objax Variables and Modules\n",
        "\n",
        "Objax Variables store values of tensors. Unlike tensors variables are mutable, i.e. the value which is stored in the variable can change.\n",
        "Since tensors are immutable, variables change their value by replacing it with new tensors.\n",
        "\n",
        "Variables are commonly used together with modules. The Module is a basic building block in Objax that stores variables and other modules. Also most modules are typically callable (i.e., implement the `__call__` method) and when called perform some computations on their variables and sub-modules.\n",
        "\n",
        "Here is an example of a simple module with one variable which performs the dot product of that variable with an input tensor:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7e7j4x0LG0kV",
        "colab_type": "code",
        "colab": {},
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772859686,
          "user_tz": 420,
          "elapsed": 525,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        }
      },
      "source": [
        "class SimpleModule(objax.Module):\n",
        "\n",
        "    def __init__(self, length):\n",
        "        self.v1 = objax.TrainVar(objax.random.normal((length,)))\n",
        "        self.v2 = objax.TrainVar(jn.ones((2,)))\n",
        "\n",
        "    def __call__(self, x):\n",
        "        return jn.dot(x, self.v1.value)\n",
        "        \n",
        "\n",
        "m = SimpleModule(3)"
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vKnFILhU7yTH",
        "colab_type": "text"
      },
      "source": [
        "Modules keep track of all variables they own, including variables in sub-modules. The `.vars()` method list all the module's variables.\n",
        "The method returns an instance of `VarCollection` which is a dictionary with several other useful methods."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KB3hOMZz8E86",
        "colab_type": "code",
        "colab": {
          "height": 204
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772860033,
          "user_tz": 420,
          "elapsed": 332,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "2dcb4565-3b03-492d-8ef3-cd34ef5386ae"
      },
      "source": [
        "module_vars = m.vars()\n",
        "\n",
        "print('type(module_vars): ', type(module_vars))\n",
        "print('isinstance(module_vars, dict): ', isinstance(module_vars, dict))\n",
        "print()\n",
        "\n",
        "print('Variable names and shapes:')\n",
        "print(module_vars)\n",
        "print()\n",
        "\n",
        "print('Variable names and values:')\n",
        "for k, v in module_vars.items():\n",
        "  print(f'{k}      {v.value}')"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "type(module_vars):  objax.variable.VarCollection\n",
            "isinstance(module_vars, dict):  True\n",
            "\n",
            "Variable names and shapes:\n",
            "(SimpleModule).v1           3 (3,)\n",
            "(SimpleModule).v2           2 (2,)\n",
            "+Total(2)                   5\n",
            "\n",
            "Variable names and values:\n",
            "(SimpleModule).v1      [-1.1010289  -0.68184537 -0.95236546]\n",
            "(SimpleModule).v2      [1. 1.]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TKQOhh6i-Oku",
        "colab_type": "text"
      },
      "source": [
        "If the `__call__` method of the module takes tensors as input and outputs tensors then it can act as a mathematical function. In the general case `__call__` can be a [multivariate](https://en.wikipedia.org/wiki/Function_(mathematics)#Multivariate_function) [vector-values](https://en.wikipedia.org/wiki/Vector-valued_function) function.\n",
        "\n",
        "The `SimpleModule` described above takes a vector of size `length` as input and outputs a scalar:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q_qjn_IG-1X9",
        "colab_type": "code",
        "colab": {
          "height": 51
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772860369,
          "user_tz": 420,
          "elapsed": 328,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "f4cf1a49-88b6-4ff5-f237-5545f9dbcaad"
      },
      "source": [
        "x = jn.ones((3,))\n",
        "y = m(x)\n",
        "print('Input: ', x)\n",
        "print('Output: ', y)"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Input:  [1. 1. 1.]\n",
            "Output:  -2.7352397\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cem5vE5GWN6V",
        "colab_type": "text"
      },
      "source": [
        "The way [jn.dot](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.dot.html) works allows us to run code on 2D tensors as well. In this case `SimpleModule` will treat the input as a batch of vectors, perform the dot product on each of them and return a vector with the results:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-ovtwEdYTKDg",
        "colab_type": "code",
        "colab": {
          "height": 136
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772860709,
          "user_tz": 420,
          "elapsed": 331,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "c9a59e31-c5aa-4212-d255-1f141eb71fbf"
      },
      "source": [
        "x = jn.array([[1., 1., 1.],\n",
        "              [1., 0., 0.],\n",
        "              [0., 1., 0.],\n",
        "              [0., 0., 1.]])\n",
        "y = m(x)\n",
        "print('Input:\\n', x)\n",
        "print('Output:\\n', y)"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Input:\n",
            " [[1. 1. 1.]\n",
            " [1. 0. 0.]\n",
            " [0. 1. 0.]\n",
            " [0. 0. 1.]]\n",
            "Output:\n",
            " [-2.7352397  -1.1010289  -0.68184537 -0.95236546]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FKbuQuERw8Oz",
        "colab_type": "text"
      },
      "source": [
        "For comparison, here is the result of calling module `m` on each row of tensor `x`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9E-uCFebwkoj",
        "colab_type": "code",
        "colab": {
          "height": 102
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772861060,
          "user_tz": 420,
          "elapsed": 329,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "b5c1189d-b6df-49e0-e50d-a8fcc2bff85c"
      },
      "source": [
        "print('Sequentially calling module on each row of 2D tensor:')\n",
        "for idx in range(x.shape[0]):\n",
        "  row_value = x[idx]\n",
        "  out_value = m(row_value)\n",
        "  print(f'm( {row_value} ) = {out_value}')"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Sequentially calling module on each row of 2D tensor:\n",
            "m( [1. 1. 1.] ) = -2.7352397441864014\n",
            "m( [1. 0. 0.] ) = -1.1010289192199707\n",
            "m( [0. 1. 0.] ) = -0.6818453669548035\n",
            "m( [0. 0. 1.] ) = -0.9523654580116272\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zhm0bt3OG7BC",
        "colab_type": "text"
      },
      "source": [
        "## How to compute gradients\n",
        "\n",
        "As shown above, modules can act as mathematical functions. It's essential in machine learning to be able to compute gradients of functions and Objax provides a simple way to do this.\n",
        "\n",
        "It's important to keep in mind that [gradients](https://en.wikipedia.org/wiki/Gradient) are usually defined for scalar-value functions, while our modules can be vector-valued. In this case we need to define additional functions which will convert vector-valued output of the module into scalar. Then we can compute gradients of scalar valued function with respect to all input variables.\n",
        "\n",
        "In the example with `SimpleModule` above let's define scalar-values loss function first:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ChsVywPkG_tD",
        "colab_type": "code",
        "colab": {
          "height": 34
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772861411,
          "user_tz": 420,
          "elapsed": 324,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "8badce00-47e4-4d7d-805c-fd0789c7a281"
      },
      "source": [
        "def loss_fn(x):\n",
        "    return m(x).sum()\n",
        "\n",
        "print('loss_fn(x) = ', loss_fn(x))"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "loss_fn(x) =  -5.4704795\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2CK6nlAkkHGY",
        "colab_type": "text"
      },
      "source": [
        "Then we create an `objax.GradValues` module which computes the gradients of `loss_fn`. We need to pass the function itself to the constructor of `objax.GradValues` as well as a `VarCollection` with the variables that `loss_fn` depends on:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K9Y6uPSYkGnO",
        "colab_type": "code",
        "colab": {},
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772861923,
          "user_tz": 420,
          "elapsed": 489,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        }
      },
      "source": [
        "# Construct a module which computes gradients\n",
        "gv = objax.GradValues(loss_fn, module_vars)"
      ],
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUK9NX_Kkif-",
        "colab_type": "text"
      },
      "source": [
        "`gv` is a module which returns the gradients of `loss_fn` and the values of `loss_fn` for the given input:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZrDVkEXlkiz9",
        "colab_type": "code",
        "colab": {
          "height": 102
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772862343,
          "user_tz": 420,
          "elapsed": 373,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "b0ceab75-c72d-456e-b8bf-1dd50eed30ef"
      },
      "source": [
        "# gv returns both gradients and values of original function\n",
        "grads, value = gv(x)\n",
        "\n",
        "print('Gradients:')\n",
        "for g, var_name in zip(grads, module_vars.keys()):\n",
        "    print(g, ' w.r.t. ', var_name)\n",
        "print()\n",
        "print('Value: ', value)"
      ],
      "execution_count": 22,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Gradients:\n",
            "[2. 2. 2.]  w.r.t.  (SimpleModule).v1\n",
            "[0. 0.]  w.r.t.  (SimpleModule).v2\n",
            "\n",
            "Value:  [DeviceArray(-5.4704795, dtype=float32)]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ohovyn4ADWCg",
        "colab_type": "text"
      },
      "source": [
        "In the example above, `grads` is a list of gradients with respect to all variables from `module_vars`. The order of gradients in the `grads` list is the same as the order of corresponding variables in `module_vars`. So `grads[0]` is the gradient of the function w.r.t. `m.v1` and `grads[1]` is the gradient w.r.t. `m.v2`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fTisDK5R_qt1",
        "colab_type": "text"
      },
      "source": [
        "## Just-in-time compilation (JIT)\n",
        "\n",
        "In the examples shown so far, the Python interpreter executes all operations one by one. This mode of execution becomes slow for larger and more complicated code.\n",
        "\n",
        "Objax provides an easy and convenient way to compile sequence of operations using `objax.Jit`: "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iR6ZRBY0NmTf",
        "colab_type": "code",
        "colab": {
          "height": 136
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772862734,
          "user_tz": 420,
          "elapsed": 358,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "713d1821-3604-46f9-d165-fa746ac4760c"
      },
      "source": [
        "jit_m = objax.Jit(m)\n",
        "y = jit_m(x)\n",
        "print('Input:\\n', x)\n",
        "print('Output:\\n', y)"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Input:\n",
            " [[1. 1. 1.]\n",
            " [1. 0. 0.]\n",
            " [0. 1. 0.]\n",
            " [0. 0. 1.]]\n",
            "Output:\n",
            " [-2.7352397  -1.1010289  -0.68184537 -0.95236546]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cH7qEQQgNpt6",
        "colab_type": "text"
      },
      "source": [
        "`Objax.Jit` can compile not only modules, but also functions and callables. In this case a variable collection should be passed to `objax.Jit`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N4YumYJxRTu7",
        "colab_type": "code",
        "colab": {
          "height": 136
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772863451,
          "user_tz": 420,
          "elapsed": 646,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "9c7eded8-8eed-4d79-a860-abdc33a6ce39"
      },
      "source": [
        "def loss_fn(x, y):\n",
        "  return ((m(x) - y) ** 2).sum()\n",
        "\n",
        "jit_loss_fn = objax.Jit(loss_fn, module_vars)\n",
        "\n",
        "x = objax.random.normal((2, 3))\n",
        "y = jn.array((-1.0, 1.0))\n",
        "\n",
        "print('x:\\n ', x)\n",
        "print('y:\\n', y)\n",
        "\n",
        "print('loss_fn(x, y): ', loss_fn(x, y))\n",
        "print('jit_loss_fn(x, y): ', jit_loss_fn(x, y))"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "x:\n",
            "  [[ 2.2491198   0.18783404  0.65321374]\n",
            " [-0.23017201  0.18411613  1.341197  ]]\n",
            "y:\n",
            " [-1.  1.]\n",
            "loss_fn(x, y):  9.577398\n",
            "jit_loss_fn(x, y):  9.577398\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KylXPvK1NzKE",
        "colab_type": "text"
      },
      "source": [
        "There is no need to use JIT if you only need to compute a single JAX operation. However JIT can give significant speedups when multiple Objax/JAX operations are chained together. The [next tutorial](Logistic_Regression.html) will show examples of how JIT is used in practice.\n",
        "\n",
        "Nevertherless the difference in execution speed with and without JIT is evident even in this simple example:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j_4S3k7ppbez",
        "colab_type": "code",
        "colab": {
          "height": 85
        },
        "executionInfo": {
          "status": "ok",
          "timestamp": 1597772867410,
          "user_tz": 420,
          "elapsed": 3935,
          "user": {
            "displayName": "Alex Kurakin",
            "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14Ggxs3OlaNvrJx7N9pYPsTha3Rztq6Jr89dSgPKY=s64",
            "userId": "14505005774567669519"
          }
        },
        "outputId": "cc226d35-7f20-4f43-d637-14b95b9ec0b0"
      },
      "source": [
        "x = objax.random.normal((100, 3))\n",
        "# gv is a module define above which compute gradients\n",
        "jit_gv = objax.Jit(gv)\n",
        "print('Timing for jit_gv:')\n",
        "%timeit jit_gv(x)\n",
        "print('Timing for gv:')\n",
        "%timeit gv(x)"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Timing for jit_gv:\n",
            "1000 loops, best of 3: 309 µs per loop\n",
            "Timing for gv:\n",
            "100 loops, best of 3: 3.57 ms per loop\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zAVF7GwPOf8n",
        "colab_type": "text"
      },
      "source": [
        "## When to use JAX and when to use Objax primitives?\n",
        "\n",
        "Attentive readers will notice that while Objax works on top of JAX, it redefines quite a few concepts from JAX. Some examples are:\n",
        "\n",
        "* `objax.GradValues` vs `jax.value_and_grad` for computing gradients.\n",
        "* `objax.Jit` vs `jax.jit` for just-in-time compilation.\n",
        "* `objax.random` vs `jax.random` to generate random numbers.\n",
        "\n",
        "All these differences originate from the fact that JAX is a stateless functional framework, while Objax provides a stateful, object-oriented way to use JAX.\n",
        "\n",
        "Mixing OOP and functional code can be quite confusing, thus we recommend to use JAX primitives only for basic mathematical operations (defined in `jax.numpy`) and use Objax primitives for everything else.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P3hsg-ShiKVf",
        "colab_type": "text"
      },
      "source": [
        "## Next: Logistic Regression Tutorial\n",
        "\n",
        "This tutorial introduces all concepts necessary to build and train a machine learning classifier. The [next tutorial](Logistic_Regression.html) shows how to apply all of them in logistic regression."
      ]
    }
  ]
}
